#!/usr/bin/env python

import sys
import sblab
import argparse
import subprocess
import os
import md5
import psycopg2

parser = argparse.ArgumentParser(description= """

DESCRIPTION:
    Execute samtools flagstat and uploads results to sblab table samtools_flagstat. In addition
    to flagstat, upload the name of input file and its md5sum.
    
    The output of samtools is also printed to stdout, so executing with --nocommit
    option is the same as straight `samtools flagstat`.
    
    Only rows not already found in samtools_flagstat are inserted so executing this script
    on the same file several times cause no change to the database.

    

EXAMPLES:
    
    flagstat_to_sblab.py el001_6.clean.8M.bam
    
REQUIRES:
    samtools on path

""", formatter_class= argparse.RawTextHelpFormatter)

parser.add_argument('bam', # nargs='?', default=sys.stdout,
                    type= str,
                    help="""Bam files to pass to samtools flagstat
                    
                    """)

parser.add_argument('--nocommit',
                    action= 'store_true',
                    help="""Do not commit changes to database.
                    """)
parser.add_argument('-v', '--verbose',
                    action= 'store_true',
                    help="""Produce verbose output: All the INSERTs to the tmp table printed to
stdout. (Note however that only rows not already in samtools_flagstat are then inserted).
                    """)

args = parser.parse_args()

# 1. Each line of output will belong to these columns. So the order of this list MUST be consistent with the output fo flagstat
# 2. These names MUST match those in table samtools_flagstat
HEADER= ['bamfile', 'md5sum', 'mate', 'tot', 'duplicates', 'mapped', 'pe_seq', 'read1', 'read2', 'prop_pe', 'itself_mate', 'singletons', 'mate_diff_chr', 'mate_diff_chr_mapq5']


# -----------------------------------------------------------------------------

def sumfile(fobj):
    '''Returns an md5 hash for an object with read() method.'''
    m = md5.new()
    while True:
        d = fobj.read(8096)
        if not d:
            break
        m.update(d)
    return(m.hexdigest())


def md5sum(fname):
    '''Returns an md5 hash for file fname, or stdin if fname is "-".'''
    if fname == '-':
        ret = sumfile(sys.stdin)
    else:
        try:
            f = file(fname, 'rb')
        except:
            return 'Failed to open file'
        ret = sumfile(f)
        f.close()
    return(ret)
# -----------------------------------------------------------------------------

cmd= 'samtools flagstat %s' %(args.bam)
p= subprocess.Popen(cmd, shell= True, stdout= subprocess.PIPE, stderr= subprocess.PIPE)
p.wait()
flagout= p.stdout.readlines()
flagerr= p.stderr.read()
if not flagerr == '':
    sys.exit(flagerr)

# --------------------------- Sanity check ------------------------------------
if not len(flagout) == 11:
    sys.exit('Unexpected number of lines returned by: %s' %(cmd))

## flagstat lines must contain these keywords:
qclist= ['in total', 'duplicates', 'mapped', 'paired in sequencing', 'read1', 'read2', 'properly paired', 'with itself and mate mapped', 'singletons', 'with mate mapped to a different chr', 'with mate mapped to a different chr (mapQ>=5)']
for qline, line in zip(qclist, flagout):
    if qline not in line:
        sys.exit('%s\nUnexpected line: "%s" (keyword "%s" not found)' %(cmd, line, qline))
# -----------------------------------------------------------------------------

flagout= [x.strip() for x in flagout]

filename= os.path.split(args.bam)[1]
mdf= md5sum(args.bam)

first_mate= [filename, mdf, '1']
second_mate= [filename, mdf, '2']
descr= []
for x in flagout:
    print(x)
    line= x.split(' ') ## Each line of output will look like "['8117297', '+', '0', 'mapped', '(100.00%:nan%)']"
    if not line[1] == '+':
        ## Small QC. Is the 2nd element the + sign?
        sys.exit('Plus sign "+" not found as second element')
    first_mate.append(int(line[0]))
    second_mate.append(int(line[2]))

pair1= {}
for k,v in zip(HEADER, first_mate):
    pair1[k]= v
pair2= {}
for k,v in zip(HEADER, second_mate):
    pair2[k]= v

# --------------------------------[ Import to sblab ]--------------------------

conn= psycopg2.connect(sblab.get_psycopgpass())
cur= conn.cursor()

sql= 'CREATE TEMP TABLE flagstat_tmp (LIKE samtools_flagstat);'
cur.execute(sql)


## Prepare for INSERT
valueinsert= ', '.join(HEADER)                    ## This will look like "bamfile, md5sum, mate,..."
headerinsert= '%(' + ')s, %('.join(HEADER) + ')s' ## This will look like "%(bamfile)s, %(md5sum)s, %(mate)s,..."
for p1, p2 in zip(pair1, pair2):
    sql= 'INSERT INTO flagstat_tmp (%s) VALUES (%s)' %(valueinsert, headerinsert)
    cur.execute(sql, pair1)
    cur.execute(sql, pair2)
    if args.verbose:
        print(cur.mogrify(sql, pair1))
        print(cur.mogrify(sql, pair2))

## Insert rows not already present.
sql= 'INSERT INTO samtools_flagstat SELECT * FROM flagstat_tmp EXCEPT SELECT * FROM samtools_flagstat; DROP TABLE flagstat_tmp;'
cur.execute(sql)

cur.close()
if not args.nocommit:
    conn.commit()
conn.close()
sys.exit()


"""
samtools flagstat el004_SLX-4526.cat.bam


3611851 + 0 in total (QC-passed reads + QC-failed reads)
0 + 0 duplicates
2774057 + 0 mapped (76.80%:nan%)
0 + 0 paired in sequencing
0 + 0 read1
0 + 0 read2
0 + 0 properly paired (nan%:nan%)
0 + 0 with itself and mate mapped
0 + 0 singletons (nan%:nan%)
0 + 0 with mate mapped to a different chr
0 + 0 with mate mapped to a different chr (mapQ>=5)

"""